"""
MRP Agent 模块

该模块实现了物料需求计划 (MRP) 的 AI Agent，能够处理用户查询并计算 MRP 计划。
主要功能包括：
1. 解析用户查询中的参数
2. 检查参数完整性
3. 调用计算工具执行 MRP 计算
4. 格式化输出结果
"""

from langchain_ollama import ChatOllama
from langchain_openai import ChatOpenAI
from langchain_core.agents import AgentFinish
from langchain_core.prompts import ChatPromptTemplate
from src.tools.mrp_calculator import calculate_mrp_tool
from src.config import OLLAMA_CONFIG, DASHSCOPE_CONFIG, SYSTEM_PROMPT, LLM_PROVIDER
import json
import re


def create_mrp_agent():
    """
    创建一个MRP Agent，可以调用MRP计算工具
    
    Returns:
        dict: 包含LLM实例、工具列表和提示词模板的字典
    """
    # 根据配置初始化LLM
    if LLM_PROVIDER == "ollama":
        llm = ChatOllama(
            base_url=OLLAMA_CONFIG["base_url"],
            model=OLLAMA_CONFIG["model"],
            temperature=OLLAMA_CONFIG["temperature"]
        )
    elif LLM_PROVIDER == "dashscope":
        llm = ChatOpenAI(
            openai_api_key=DASHSCOPE_CONFIG["api_key"],
            openai_api_base=DASHSCOPE_CONFIG["base_url"],
            model=DASHSCOPE_CONFIG["model"],
            temperature=DASHSCOPE_CONFIG["temperature"]
        )
        print('3333')
    else:
        raise ValueError(f"不支持的LLM提供商: {LLM_PROVIDER}")

    # 定义工具列表
    tools = [calculate_mrp_tool]

    # 创建提示词模板
    prompt = ChatPromptTemplate.from_messages([
        ("system", SYSTEM_PROMPT),
        ("human", "{input}"),
    ])

    # 简单实现：直接调用工具，不使用复杂的Agent
    return {
        "llm": llm,
        "tools": tools,
        "prompt": prompt
    }


def extract_mrp_params(query: str) -> dict:
    """
    从查询中提取MRP参数
    
    Args:
        query (str): 用户查询字符串
        
    Returns:
        dict: 提取到的参数字典
    """
    params = {}
    try:
        total_demand_match = re.search(r"总需求量:\s*([\d.]+|未提供)", query)
        if total_demand_match and total_demand_match.group(1) != "未提供":
            params["total_demand"] = float(total_demand_match.group(1))
            
        daily_demand_match = re.search(r"日需求量:\s*([\d.]+|未提供)", query)
        if daily_demand_match and daily_demand_match.group(1) != "未提供":
            params["daily_demand"] = float(daily_demand_match.group(1))
            
        order_qty_match = re.search(r"订货批量:\s*([\d.]+|未提供)", query)
        if order_qty_match and order_qty_match.group(1) != "未提供":
            params["order_qty"] = float(order_qty_match.group(1))
            
        reorder_point_match = re.search(r"订货点:\s*([\d.]+|未提供)", query)
        if reorder_point_match and reorder_point_match.group(1) != "未提供":
            params["reorder_point"] = float(reorder_point_match.group(1))
            
        lead_time_days_match = re.search(r"订货提前期:\s*(\d+|未提供)", query)
        if lead_time_days_match and lead_time_days_match.group(1) != "未提供":
            params["lead_time_days"] = int(lead_time_days_match.group(1))
            
        start_date_match = re.search(r"生产开始日期:\s*([\d-]+|未提供)", query)
        if start_date_match and start_date_match.group(1) != "未提供":
            params["start_date"] = start_date_match.group(1)
            
        initial_stock_match = re.search(r"初始库存:\s*([\d.]+|未提供)", query)
        if initial_stock_match and initial_stock_match.group(1) != "未提供":
            params["initial_stock"] = float(initial_stock_match.group(1))
    except Exception:
        pass
    
    return params


def check_missing_params(params: dict) -> list:
    """
    检查缺失的参数
    
    Args:
        params (dict): 已提供的参数字典
        
    Returns:
        list: 缺失参数的列表
    """
    required_params = [
        "total_demand", "daily_demand", "order_qty", 
        "reorder_point", "lead_time_days", "start_date", "initial_stock"
    ]
    missing_params = []
    
    for param in required_params:
        if param not in params:
            missing_params.append(param)
    
    return missing_params


def run_mrp_agent(query: str):
    """
    运行MRP Agent来处理查询
    
    Args:
        query (str): 用户查询字符串
        
    Yields:
        str: 计算结果的各个部分，以流式方式返回
    """
    # 检查查询中是否包含计算MRP的关键词
    if "计算MRP" in query or "计算mrp" in query or "MRP计划" in query:
        # 解析参数
        try:
            total_demand = float(re.search(r"总需求量:\s*([\d.]+)", query).group(1))
            daily_demand = float(re.search(r"日需求量:\s*([\d.]+)", query).group(1))
            order_qty = float(re.search(r"订货批量:\s*([\d.]+)", query).group(1))
            reorder_point = float(re.search(r"订货点:\s*([\d.]+)", query).group(1))
            lead_time_days = int(re.search(r"订货提前期:\s*(\d+)", query).group(1))
            start_date = re.search(r"生产开始日期:\s*([\d-]+)", query).group(1)
            initial_stock = float(re.search(r"初始库存:\s*([\d.]+)", query).group(1))
            
            # 调用计算工具
            orders_info, daily_stock_plan = calculate_mrp_tool.invoke({
                "total_demand": total_demand,
                "daily_demand": daily_demand,
                "order_qty": order_qty,
                "reorder_point": reorder_point,
                "lead_time_days": lead_time_days,
                "start_date": start_date,
                "initial_stock": initial_stock
            })
            
            # 格式化输出结果
            yield "MRP计算结果:\n\n"
            
            # 输出订单信息表
            yield "订单与到货信息表:\n"
            yield "| 序号 | 订单下达日期 | 到货日期 | 到货后库存 |\n"
            yield "|------|-------------|----------|-----------|\n"
            for order in orders_info:
                yield f"| {order['序号']} | {order['order_date']} | {order['arrival_date']} | {order['stock_after_arrival']}吨 |\n"
            
            yield "\n每日库存计划:\n"
            yield "| 日期 | 库存(日初始库存) | 事件 |\n"
            yield "|------|----------------|------|\n"
            for day in daily_stock_plan:
                yield f"| {day['date']} | {day['start_stock']}吨 | {day['events']} |\n"
        except Exception as e:
            yield f"解析参数或计算过程中出现错误: {str(e)}\n"
            yield "请确保输入格式正确，包含所有必要参数。\n"
    else:
        # 简化实现：直接调用LLM处理查询
        agent = create_mrp_agent()
        llm = agent["llm"]
        prompt = agent["prompt"]
        
        # 格式化提示词
        formatted_prompt = prompt.format_messages(input=query)
        # 修改为流式输出
        for chunk in llm.stream(formatted_prompt):
            yield chunk.content


def run_interactive_mrp_agent(query: str, provided_params: dict = None):
    """
    运行交互式MRP Agent来处理查询
    
    Args:
        query (str): 用户查询字符串
        provided_params (dict, optional): 已提供的参数字典
        
    Returns:
        dict: 包含响应内容和完成状态的字典
    """
    if provided_params is None:
        provided_params = {}
    
    # 从查询中提取已有的参数
    extracted_params = extract_mrp_params(query)
    
    # 合并参数，provided_params优先级更高
    params = {**extracted_params, **provided_params}
    
    # 检查缺失的参数
    missing_params = check_missing_params(params)
    
    if missing_params:
        # 有缺失参数，将问题交给LLM进行一轮对话
        agent = create_mrp_agent()
        llm = agent["llm"]
        prompt = agent["prompt"]
        
        # 构造提示语，让LLM引导用户提供完整信息
        prompt_text = f"""
用户想要计算MRP计划，但提供了不完整的信息。以下是用户已经提供的信息：
{query}

但是还需要以下参数才能进行计算：
1. 总需求量（吨）
2. 日需求量（吨）
3. 订货批量（吨）
4. 订货点（吨）
5. 订货提前期（天）
6. 生产开始日期（格式 YYYY-MM-DD）
7. 初始库存（吨）

请以友好的方式询问用户提供这些缺失的信息，并解释为什么需要这些信息。
        """
        
        formatted_prompt = prompt.format_messages(input=prompt_text)
        response = ""
        for chunk in llm.stream(formatted_prompt):
            response += chunk.content
            
        return {
            "response": response,
            "completed": False
        }
    else:
        # 参数齐全，执行计算
        try:
            # 调用计算工具
            orders_info, daily_stock_plan = calculate_mrp_tool.invoke(params)
            
            # 格式化输出结果
            result = "MRP计算结果:\n\n"
            
            # 输出订单信息表
            result += "订单与到货信息表:\n"
            result += "| 序号 | 订单下达日期 | 到货日期 | 到货后库存 |\n"
            result += "|------|-------------|----------|-----------|\n"
            for order in orders_info:
                result += f"| {order['序号']} | {order['order_date']} | {order['arrival_date']} | {order['stock_after_arrival']}吨 |\n"
            
            result += "\n每日库存计划:\n"
            result += "| 日期 | 库存(日初始库存) | 事件 |\n"
            result += "|------|----------------|------|\n"
            for day in daily_stock_plan:
                result += f"| {day['date']} | {day['start_stock']}吨 | {day['events']} |\n"
                
            return {
                "response": result,
                "completed": True
            }
        except Exception as e:
            return {
                "response": f"计算过程中出现错误: {str(e)}",
                "completed": True
            }